"""
calulate rainfall spatial and temporal coefficient of variability
"""
# -*- coding: utf-8 -*-
import os
import numpy as np
import pandas as pd


def get_cum_rainfall(dat_dir):
    """
    calculate cumulative rainfall using files in dat_dir
    """
    files = os.listdir(dat_dir)
    cum_rainfall = pd.DataFrame(np.zeros((90, 90)))
    for filename in files:
        filepath = os.path.join(dat_dir, filename)
        dat = pd.read_csv(filepath, index_col=False, header=None)
        cum_rainfall = cum_rainfall + dat
    return cum_rainfall


def calc_spatial_var(cum_rainfall):
    """calculate spatial variability using cum_rainfall DataFrame"""
    average = cum_rainfall.stack().mean()

    cv_spatial = np.sqrt((cum_rainfall / average - 1).pow(2).stack().mean())
    return cv_spatial


def get_areal_average(dat_dir):
    """
    calculate areal average of each file in dat_dir
    return a list of areal_average
    """
    files = os.listdir(dat_dir)
    areal_average = []
    for filename in files:
        filepath = os.path.join(dat_dir, filename)
        dat = pd.read_csv(filepath, index_col=False, header=None)
        average = dat.stack().mean()
        areal_average.append(average)
    return areal_average


def calc_temporal_var(areal_average):
    """calculate temporal variability using areal average rainfall"""
    areal_average = np.asarray(areal_average)
    average = np.mean(areal_average)
    cv_temporal = np.sqrt(np.mean(np.power((areal_average / average - 1), 2)))
    return cv_temporal


def main():
    """main function"""
    dat_dir = r"F:/research/rainfall_estimation/dat/nimrod_output/2008/R8"
    cum_rainfall = get_cum_rainfall(dat_dir)
    # write accumulative rainfall
    cum_rainfall_filepath = os.path.join(
        r"F:/research/rainfall_estimation/dat/nimrod_output/2008/cum_rainfall", "R8.asc")
    cum_rainfall.to_csv(cum_rainfall_filepath, index=False, header=False)
    # calculate spatial variation
    cv_spatial = calc_spatial_var(cum_rainfall)
    # calculate temporal variation
    areal_average = get_areal_average(dat_dir)
    cv_temporal = calc_temporal_var(areal_average)
    print("spatial cv: ", cv_spatial)
    print("temporal cv: ", cv_temporal)


if __name__ == '__main__':
    main()
